#include <iostream>
#include <unistd.h>
#include <vector>
#include "task.hpp"

using namespace std;

// 枚举错误类型
enum Error
{
    UsageError = 1,
    ArgError,
    PipeError
};

class Channel
{
public:
    Channel(int wfd, pid_t pid, const string &name)
        : _wfd(wfd), _pid(pid), _pipe_name(name)
    {
    }
    string name() { return _pipe_name; }
    int wfd() { return _wfd; }
    pid_t pid() { return _pid; }

    ~Channel()
    {
    }

    void PrintDebug()
    {
        cout << "channel name: " << _pipe_name << endl;
        cout << "sub process id: " << _pid << endl;
        cout << "write side: " << _wfd << endl;
    }

    void Close() { close(_wfd); };
private:
    // 连接的子进程
    pid_t _pid;
    string _pipe_name;
    int _wfd;
};

class ProcessPool
{
public:
    ProcessPool(int sub_num_process)
        : _sub_num_process(sub_num_process)
    {
    }
    int CreateProcess(work_t worker)
    {
        // 创建_sub_num_process个子进程并建立管道
        for (int i = 0; i < _sub_num_process; i++)
        {
            // 创建信道
            int pipefd[2]{0};
            int n = pipe(pipefd);
            // 创建失败
            if (n < 0)
                return PipeError;
            pid_t id = fork();
            if (id == 0)
            {
                // 子
                // 我们需要子进程关闭写端，只要读端
                close(pipefd[1]);
                // 执行任务
                // 把读端和标准输入连接,这样就可以通过标准输入直接读到任务
                dup2(pipefd[0], 0);
                worker();
                exit(0);
            }
            // 父
            // 父进程关闭读端
            string name = "channel-" + to_string(i);
            close(pipefd[0]);
            // 管理信道
            _channels.push_back(Channel(pipefd[1], id, name));
        }
        return 0;
    }

    void KillAll()
    {
        for(auto &channel : _channels)
        {
            channel.Close();
            cout << channel.name() << " close done" << ", sub process quit now: " << channel.pid() << endl;
            pid_t rid = waitpid(channel.pid(), nullptr, 0);
            if(rid == channel.pid())
            {
                cout << "wait sub process : " << rid << " success!" << endl;
            }

        }
    }

    // void Wait()
    // {
    //     for(auto &channel: _channels)
    //     {
    //         pid_t pid = channel.pid();
    //         pid_t rid = waitpid(pid, nullptr, 0);
    //         if(rid == pid)
    //         {
    //             cout << "wait sub process : " << pid << " success!" << endl;
    //         }
    //     }
    // }
    // 实现负载均衡的选择一个管道，让那个进程工作，实际上就是向哪个管道写入工作码
    // 返回该管道的索引
    int NextChannel()
    {
        // 每次按顺序取当前索引下的任务，再令索引++，实现轮询操作
        static int next = 0;
        int c = next;
        next++;
        next %= _channels.size();
        return c;
    }

    // 分配任务
    void SendTask(int index, uint32_t work_code)
    {
        cout << "send code: " << work_code << ' ' << _channels[index].name() << " sub process id: " << _channels[index].pid() << endl;
        write(_channels[index].wfd(), &work_code, sizeof(work_code));
    }

    // 输出进程池信息
    void Debug()
    {
        for (auto &e : _channels)
        {
            e.PrintDebug();
        }
    }

    ~ProcessPool()
    {
    }

private:
    // 进程数量
    int _sub_num_process;
    // 管道
    vector<Channel> _channels;
};

void Usage(const string& name)
{
    cout << "the correct use is like " << name << " number(1, 2, ...)" << endl;
}

void CtrlProcessPool(ProcessPool *process_pool_ptr, int n)
{
    while (n--)
    {
        // 选择一个管道
        // 实现负载均衡
        int channel = process_pool_ptr->NextChannel();
        //获取任务
        uint32_t work_code = NextTask();
        //发送任务
        process_pool_ptr->SendTask(channel, work_code);
        sleep(1);
    }
}
int main(int argc, char *argv[])
{
    // 如果用法错了
    if (argc != 2)
    {
        Usage(argv[0]);
        return UsageError;
    }
    // 获取创建进程池初始化数目
    int sub_num_process = stoi(argv[1]);
    if (sub_num_process <= 0)
        return ArgError;
    // 设置随机数种子
    srand((uint64_t)time(nullptr));
    // 1.创建进程池
    ProcessPool *process_pool_ptr = new ProcessPool(sub_num_process);
    process_pool_ptr->CreateProcess(Worker);
    // 2.控制子进程
    CtrlProcessPool(process_pool_ptr, 10);
    //回收子进程
    process_pool_ptr->KillAll();
    //process_pool_ptr->Wait();
    delete process_pool_ptr;
    return 0;
}